{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hgvFF5YOfuJW"
      },
      "outputs": [],
      "source": [
        "# Copyright 2023 The TensorFlow GNN Authors.\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vV5wjlQwf9Ky"
      },
      "source": [
        "\n",
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/gnn/blob/master/examples/tutorials/kdd_2023/code_tutorial_datasets.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/gnn/blob/main/examples/tutorials/kdd_2023/code_tutorial_datasets.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8CJe-bK4eIrH"
      },
      "source": [
        "# Welcome\n",
        "The goal of this tutorial is to show how to define a new in-memory dataset in TF-GNN."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vtyMTRkozVKx"
      },
      "source": [
        "## Imports"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "myGdQat5esDh"
      },
      "source": [
        "At first, you need to install TF-GNN and import required libraries. As TF-GNN is actively developped, we include code that is on github but **not yet in pip package** in a hidden cell."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DuMeO0jVfQiC"
      },
      "outputs": [],
      "source": [
        "!pip install tensorflow_gnn~=0.6.0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Dse3L5crzAkZ"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from absl import app\n",
        "from absl import flags\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import tensorflow_gnn as tfgnn\n",
        "from tensorflow_gnn import runner\n",
        "# from tensorflow_gnn.experimental.in_memory import data_providers  # In hidden-cell below\n",
        "from tensorflow_gnn.models import vanilla_mpnn\n",
        "from tensorflow_gnn.experimental.in_memory import datasets\n",
        "from tensorflow_gnn.models import gcn"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SuBd-HNGAq63"
      },
      "source": [
        "### Hidden Cell (data_providers)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mFLDYqf9ApnM"
      },
      "outputs": [],
      "source": [
        "\n",
        "# Copyright 2023 The TensorFlow GNN Authors. All Rights Reserved.\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     http://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License.\n",
        "# ==============================================================================\n",
        "\"\"\"Creates `DatasetProvider` and `Tasks` instances to invoke TF-GNN runner.\"\"\"\n",
        "\n",
        "from __future__ import annotations\n",
        "import abc\n",
        "import dataclasses\n",
        "import functools\n",
        "from typing import Any, Callable\n",
        "\n",
        "import tensorflow as tf\n",
        "from tensorflow_gnn import runner\n",
        "from tensorflow_gnn.experimental.in_memory import datasets\n",
        "from tensorflow_gnn.experimental.in_memory import int_arithmetic_sampler as ia_sampler\n",
        "from tensorflow_gnn.experimental.sampler import link_samplers\n",
        "from tensorflow_gnn.experimental.sampler import subgraph_pipeline\n",
        "from tensorflow_gnn.graph import adjacency as adj\n",
        "from tensorflow_gnn.graph import graph_constants\n",
        "from tensorflow_gnn.graph import graph_tensor as gt\n",
        "from tensorflow_gnn.proto import graph_schema_pb2\n",
        "from tensorflow_gnn.sampler import sampling_spec_builder\n",
        "from tensorflow_gnn.sampler import sampling_spec_pb2\n",
        "\n",
        "\n",
        "RANDOM_UNIFORM = sampling_spec_builder.SamplingStrategy.RANDOM_UNIFORM\n",
        "\n",
        "\n",
        "class SizedDatasetProvider(runner.DatasetProvider):\n",
        "\n",
        "  @property\n",
        "  @abc.abstractmethod\n",
        "  def cardinality(self) -\u003e int:\n",
        "    raise NotImplementedError()\n",
        "\n",
        "\n",
        "@dataclasses.dataclass(frozen=True)\n",
        "class TaskedData:\n",
        "  \"\"\"Bundles instances required to train by `tfgnn.runner`: dataset and task.\"\"\"\n",
        "  task: runner.Task\n",
        "  provider: SizedDatasetProvider\n",
        "  validation_provider: None|SizedDatasetProvider = None\n",
        "\n",
        "\n",
        "class Sampling:\n",
        "  \"\"\"Union (one-of) sampling options.\"\"\"\n",
        "\n",
        "  @property\n",
        "  def spec(self) -\u003e None | sampling_spec_pb2.SamplingSpec:\n",
        "    return self._spec\n",
        "\n",
        "  @property\n",
        "  def sample_sizes(self) -\u003e None | tuple[int, ...]:\n",
        "    return self._sample_sizes\n",
        "\n",
        "  def __init__(self, spec: None | sampling_spec_pb2.SamplingSpec,\n",
        "               sample_sizes: None | tuple[int, ...]):\n",
        "    if int(spec is None) + int(sample_sizes is None) != 1:\n",
        "      raise ValueError('Expecting exactly one of `spec` and `sample_sizes`')\n",
        "    self._spec = spec\n",
        "    self._sample_sizes = sample_sizes\n",
        "\n",
        "  @staticmethod\n",
        "  def from_spec(spec: sampling_spec_pb2.SamplingSpec) -\u003e Sampling:\n",
        "    return Sampling(spec=spec, sample_sizes=None)\n",
        "\n",
        "  @staticmethod\n",
        "  def from_sample_sizes(sample_sizes: list[int]) -\u003e Sampling:\n",
        "    return Sampling(spec=None, sample_sizes=tuple(sample_sizes))\n",
        "\n",
        "\n",
        "Sampling.SMALL_SAMPLING = Sampling.from_sample_sizes([5, 5])\n",
        "\n",
        "\n",
        "def provide_link_prediction_data(\n",
        "    *,\n",
        "    dataset_name: None|str = None,\n",
        "    graph_data: None|datasets.LinkPredictionGraphData = None,\n",
        "    source_sampling: None|Sampling = None,\n",
        "    target_sampling: None|Sampling = None,\n",
        "    task: None|runner.Task = None,\n",
        "    ) -\u003e TaskedData:\n",
        "  \"\"\"Provides `TaskedData` for link prediction tasks.\n",
        "\n",
        "  Given a graph (indicated by `dataset_name` or `graph_data`), a dataset\n",
        "  iterating all edges and producing subgraph `GraphTensor` around each edge will\n",
        "  created. The size of all `GraphTensor`s (depth and fanout) is determined by\n",
        "  `sampling`. Positive edges are sampled uniformly. For each positive edge\n",
        "  connecting `source-\u003etarget`, negative edges `source-\u003en_target` and\n",
        "  `n_source-\u003etarget` will be sampled where `n_source` and `n_target` are sampled\n",
        "  uniformly at random.\n",
        "\n",
        "  Args:\n",
        "    dataset_name: Mutually-exclusive with `graph_data`. If given, it must be\n",
        "      name of node link prediction dataset recognized by in_memory/datasets.py,\n",
        "      e.g., \"ogbl-collab\". This is used to create `graph_data` without any\n",
        "      modifications (edges will not be undirected, self-connections will not be\n",
        "      added, etc).\n",
        "    graph_data: Mutually-exclusive with `dataset_name`. It provides the list of\n",
        "      (positive) training edges and validation (positive and negative) edges.\n",
        "    source_sampling: Specifies how each `GraphTensor` instance will be sampled\n",
        "      around an (positive or negative) endpoint `source` around edge\n",
        "      `source-\u003etarget`. If None, defaults to `SMALL_SAMPLING` (samples 5 nodes\n",
        "      from each edge-set connected to `source`, and for each sampled node,\n",
        "      samples 5 of its neighbors per edgeset).\n",
        "    target_sampling: Similar to `source_sampling`, but for `target` endpoint.\n",
        "    task: Runner task to apply. If not given, defaults to\n",
        "      `DotProductLinkPrediction`. If given, user must provide a task that works\n",
        "      well with the produced subgraphs.\n",
        "\n",
        "  Returns:\n",
        "    `TaskedData` (`td`) with `td.task` set to a link-prediction task and\n",
        "    `provider` set to `DatasetProvider` that yields a dataset of `GraphTensor`\n",
        "    instances. Each instance is a subgraph sample around source-node of an edge\n",
        "    unioned with a subgraph sample around target-node of the same edge. The\n",
        "    positions of the edge endpoints are specified in edge-sets `_readout/source`\n",
        "    and `_readout/target` and the label of the edge (positive or negative) can\n",
        "    accessed as `graph.node_sets['_readout']['label']`.\n",
        "  \"\"\"\n",
        "  graph_data = _resolve_graph_data(\n",
        "      datasets.LinkPredictionGraphData, dataset_name, graph_data)\n",
        "  assert isinstance(graph_data, datasets.LinkPredictionGraphData)\n",
        "  schema = graph_data.graph_schema()\n",
        "\n",
        "  providers = _provide_link_prediction_data_with_int_arithmetic_sampling(\n",
        "      graph_data,\n",
        "      _resolve_sampling_spec(\n",
        "          schema, graph_data.source_node_set_name, source_sampling),\n",
        "      _resolve_sampling_spec(\n",
        "          schema, graph_data.target_node_set_name, target_sampling))\n",
        "  train_provider, validation_provider = providers\n",
        "\n",
        "  if task is None:\n",
        "    task = runner.DotProductLinkPrediction()\n",
        "\n",
        "  return TaskedData(task=task, provider=train_provider,\n",
        "                    validation_provider=validation_provider)\n",
        "\n",
        "\n",
        "def provide_node_classification_data(\n",
        "    *,\n",
        "    dataset_name: None|str = None,\n",
        "    graph_data: None|datasets.NodeClassificationGraphData = None,\n",
        "    sampling: None|Sampling = None) -\u003e TaskedData:\n",
        "  \"\"\"Provides `TaskedData` for node classification tasks.\"\"\"\n",
        "  graph_data = _resolve_graph_data(\n",
        "      datasets.NodeClassificationGraphData, dataset_name, graph_data)\n",
        "  assert isinstance(graph_data, datasets.NodeClassificationGraphData)\n",
        "\n",
        "  sampler = ia_sampler.GraphSampler(\n",
        "      graph_data, sampling_mode=ia_sampler.EdgeSampling.WITH_REPLACEMENT)\n",
        "\n",
        "  node_split = graph_data.node_split()\n",
        "  schema = graph_data.graph_schema()\n",
        "  provider_kwargs = dict(\n",
        "      sampling_pipeline=subgraph_pipeline.SamplingPipeline(\n",
        "          schema,\n",
        "          _resolve_sampling_spec(schema, graph_data.labeled_nodeset, sampling),\n",
        "          sampler.make_edge_sampler,\n",
        "          create_node_features_lookup_factory(graph_data, True)),\n",
        "      labeled_node_set_name=graph_data.labeled_nodeset)\n",
        "  train_provider = _NodeClassificationSamplingProvider(\n",
        "      node_split.train.shape[0],\n",
        "      dataset_fn=functools.partial(\n",
        "          _create_nodes_and_labels_dataset_from_eager_tensors,\n",
        "          node_split.train, graph_data.labels(), shuffle=True),\n",
        "      **provider_kwargs)\n",
        "  validation_provider = _NodeClassificationSamplingProvider(\n",
        "      node_split.validation.shape[0],\n",
        "      dataset_fn=functools.partial(\n",
        "          _create_nodes_and_labels_dataset_from_eager_tensors,\n",
        "          node_split.validation, graph_data.labels()),\n",
        "      **provider_kwargs)\n",
        "  task = _ManyNodesClassification(\n",
        "      graph_data.labeled_nodeset, label_fn=_readout_seed_node_labels,\n",
        "      num_classes=graph_data.num_classes())\n",
        "  return TaskedData(\n",
        "      task=task, provider=train_provider,\n",
        "      validation_provider=validation_provider)\n",
        "\n",
        "\n",
        "## \u003cHelper functions for Node Classification\u003e\n",
        "\n",
        "\n",
        "def _create_nodes_and_labels_dataset_from_eager_tensors(\n",
        "    seed_nodes: tf.Tensor,\n",
        "    labels: tf.Tensor,\n",
        "    context: None|tf.distribute.InputContext = None,\n",
        "    *,\n",
        "    shuffle=False) -\u003e tf.data.Dataset:\n",
        "  \"\"\"Creates dataset of `(node, label)` from `seed_nodes` and `labels`.\"\"\"\n",
        "  context = context or tf.distribute.InputContext()\n",
        "  dataset = (tf.data.Dataset.from_tensor_slices(seed_nodes)\n",
        "             .shard(context.num_input_pipelines, context.input_pipeline_id))\n",
        "  if shuffle:\n",
        "    dataset = dataset.shuffle(\n",
        "        seed_nodes.shape[0], reshuffle_each_iteration=True)\n",
        "  dataset = dataset.map(lambda node: (node, tf.gather(labels, node)))\n",
        "  return dataset\n",
        "\n",
        "\n",
        "def _readout_seed_node_labels(graph: gt.GraphTensor) -\u003e tuple[\n",
        "    gt.GraphTensor, graph_constants.Field]:\n",
        "  labels = tf.gather(\n",
        "      graph.node_sets['_readout']['label'],\n",
        "      graph.edge_sets['_readout/seed'].adjacency.target)\n",
        "  return graph, labels\n",
        "\n",
        "\n",
        "class _ManyNodesClassification(runner.RootNodeMulticlassClassification):\n",
        "  \"\"\"Classification setup with possibly multiple labeled nodes.\"\"\"\n",
        "\n",
        "  def gather_activations(self, graph: gt.GraphTensor) -\u003e graph_constants.Field:\n",
        "    x = graph.node_sets[self._node_set_name][graph_constants.HIDDEN_STATE]\n",
        "    return tf.gather(x, graph.edge_sets['_readout/seed'].adjacency.source)\n",
        "\n",
        "\n",
        "class _NodeClassificationSamplingProvider(SizedDatasetProvider):\n",
        "  \"\"\"DatasetProvider for Node Classification.\"\"\"\n",
        "\n",
        "  def __init__(\n",
        "      self,\n",
        "      cardinality: int,\n",
        "      dataset_fn: Callable[[tf.distribute.InputContext], tf.data.Dataset],\n",
        "      sampling_pipeline: subgraph_pipeline.SamplingPipeline,\n",
        "      labeled_node_set_name: graph_constants.NodeSetName,\n",
        "      sampling_batch_size: int = 100):\n",
        "    \"\"\"Constructs DatasetProvider for Node Classification.\n",
        "\n",
        "    Args:\n",
        "      cardinality: Total number of examples in (one iteration of) dataset.\n",
        "      dataset_fn: Callable returns dataset that with each element containing\n",
        "        `(seed ID, label)`. The first will be batched to invoke\n",
        "        `sampling_pipeline`.\n",
        "      sampling_pipeline: Maps seed IDs from dataset into `GraphTensor`s.\n",
        "      labeled_node_set_name: The node set that the labels correspond to.\n",
        "      sampling_batch_size: Batch size to send to `sampling_model`.\n",
        "    \"\"\"\n",
        "    super().__init__()\n",
        "    self._cardinality = cardinality\n",
        "    self._dataset_fn = dataset_fn\n",
        "    self._sampling_pipeline = sampling_pipeline\n",
        "    self._node_set_name = labeled_node_set_name\n",
        "    self._batch_size = sampling_batch_size\n",
        "\n",
        "  @property\n",
        "  def cardinality(self) -\u003e int:\n",
        "    return self._cardinality\n",
        "\n",
        "  def get_dataset(self, context: tf.distribute.InputContext) -\u003e tf.data.Dataset:\n",
        "    \"\"\"Returns dataset with `GraphTensor` elements for node classification.\n",
        "\n",
        "    Each `GraphTensor` instance is to classify one or more nodes. Specifically,\n",
        "    edge-set `'_readout/seed'` will contain `K` edges, with source indices\n",
        "    pointing to labeled nodes. Further, node-set `'_readout'` will contain `K`\n",
        "    nodes, with feature `'label'` containing the corresponding node labels,\n",
        "    where `K` is the number of labeled nodes in `GraphTensor`.\n",
        "\n",
        "    Args:\n",
        "      context: To select the data shard. If constructed with default constructor\n",
        "        (`tf.distribute.InputContext()`), then all data will used by worker.\n",
        "    \"\"\"\n",
        "    sampler = self._sampling_pipeline  # for short.\n",
        "    return (\n",
        "        self._dataset_fn(context)\n",
        "        .batch(self._batch_size)  # For invoking sampler (unbatch to follow).\n",
        "        # Sample subgraph.\n",
        "        .map(lambda nodes, labels: (sampler(nodes), nodes, labels),\n",
        "             num_parallel_calls=tf.data.AUTOTUNE, deterministic=False)\n",
        "        .unbatch()\n",
        "        .map(self._add_structured_readout,\n",
        "             num_parallel_calls=tf.data.AUTOTUNE, deterministic=False)\n",
        "    )\n",
        "\n",
        "  def _add_structured_readout(self, graph, node_id, label):\n",
        "    node_sets = dict(graph.node_sets)\n",
        "    edge_sets = dict(graph.edge_sets)\n",
        "    ones = tf.ones([1], dtype=graph.indices_dtype)\n",
        "    node_sets['_readout'] = gt.NodeSet.from_fields(\n",
        "        sizes=ones, features={'label': tf.expand_dims(label, 0)})\n",
        "    node_ids = tf.cast(\n",
        "        graph.node_sets[self._node_set_name]['#id'], node_id.dtype)\n",
        "    node_id_position = tf.argmax(node_id == node_ids,\n",
        "                                 output_type=graph.indices_dtype)\n",
        "    edge_sets['_readout/seed'] = gt.EdgeSet.from_fields(\n",
        "        sizes=ones, adjacency=adj.Adjacency.from_indices(\n",
        "            source=(self._node_set_name, tf.expand_dims(node_id_position, 0)),\n",
        "            target=('_readout', tf.zeros([1], dtype=graph.indices_dtype)),\n",
        "        ))\n",
        "    return gt.GraphTensor.from_pieces(\n",
        "        context=graph.context, edge_sets=edge_sets, node_sets=node_sets)\n",
        "\n",
        "\n",
        "## \u003cHelper functions for Link Prediction.\u003e\n",
        "\n",
        "\n",
        "def _create_labeled_edges_dataset_from_eager_tensors(\n",
        "    sources: tf.Tensor,\n",
        "    targets: tf.Tensor,\n",
        "    context: None|tf.distribute.InputContext = None,\n",
        "    *,\n",
        "    shuffle=False,\n",
        "    negative_links_sampling: NegativeLinksSampling,\n",
        "    ) -\u003e tf.data.Dataset:\n",
        "  \"\"\"Creates link-prediction dataset from `sources`, `targets`, and negatives.\n",
        "\n",
        "  It must be that `source_nodes` and `target_nodes` are both vectors with equal\n",
        "  shape, with each `(source_nodes[i], target_nodes[i])` indicating a positive\n",
        "  edge.\n",
        "\n",
        "  For each positive edge, a total of\n",
        "  ```num_negatives_per_source + num_negatives_per_target```\n",
        "  negative examples will be sampled (uniformly at random).\n",
        "\n",
        "  Args:\n",
        "    sources: Eager vector of size M. It is (currently) assumed that these\n",
        "      nodes are integral, and come from range `[0, max_source_node]`. Each\n",
        "      `(sources[i], targets[i])` indicate an edge.\n",
        "    targets: Eager vector of size M. It is (currently) assumed that these\n",
        "      nodes are integral, and come from range `[0, max_target_node]`.\n",
        "    context: If given, a subset of the data will be used that is designated for\n",
        "      to the worker by the `context`.\n",
        "    shuffle: If set, input edges (`sources, targets`) will be shuffled.\n",
        "    negative_links_sampling: Contains functions that generate negative source\n",
        "      and target node IDs.\n",
        "\n",
        "  Returns:\n",
        "    tf.data.Dataset with total examples:\n",
        "    `pe_matrix.shape[0] *(1+num_negatives_per_source+num_negatives_per_target)`.\n",
        "    Each positive edge will be followed by its negatives. The returned dataset\n",
        "    will yield pair: `(edge, label)`, where `label` is binary tf.float32, and\n",
        "    `edge` is a 2-size vector containing `[src, tgt]`: the vector can be a row\n",
        "    of `positive_edge_matrix` (in which case, `label == 1.0`), or one endpoint\n",
        "    is from `[src, tgt]` and the other is (uniformly) randomly sampled.\n",
        "  \"\"\"\n",
        "  if (len(sources.shape) != 1 or len(targets.shape) != 1\n",
        "      or int(sources.shape[0]) != targets.shape[0]):\n",
        "    raise ValueError(\n",
        "        'Expecting vectors `sources` and `targets` of same size')\n",
        "\n",
        "  context = context or tf.distribute.InputContext()\n",
        "  ds = (\n",
        "      tf.data.Dataset.range(targets.shape[0])\n",
        "      .shard(context.num_input_pipelines, context.input_pipeline_id))\n",
        "\n",
        "  if shuffle:\n",
        "    ds = ds.shuffle(targets.shape[0], reshuffle_each_iteration=True)\n",
        "\n",
        "  return amend_negative_edges(\n",
        "      ds.map(lambda i: (tf.gather(sources, i), tf.gather(targets, i))),\n",
        "      negative_links_sampling=negative_links_sampling)\n",
        "\n",
        "\n",
        "@dataclasses.dataclass(frozen=True)\n",
        "class NegativeLinksSampling:\n",
        "  \"\"\"Instructs how to sample negative edges for link-prediction task.\"\"\"\n",
        "\n",
        "  # Function that can generate negative source node IDs. `int` argument\n",
        "  # indicates the number of desired negative IDs.\n",
        "  negative_sources_fn: Callable[[int], tf.Tensor]\n",
        "  # Function that can generate negative target node IDs. `int` argument\n",
        "  # indicates the number of desired negative IDs.\n",
        "  negative_targets_fn: Callable[[int], tf.Tensor]\n",
        "\n",
        "  # For each positive edge `(src, tgt)`, the `src` endpoint will be paired with\n",
        "  # this many negative targets.\n",
        "  num_negatives_per_source: int\n",
        "  # ... and the `tgt` endpoint will be paired with this many negative sources.\n",
        "  num_negatives_per_target: int\n",
        "\n",
        "  def generate_sources(self) -\u003e tf.Tensor:\n",
        "    \"\"\"Returns negative-sampled sources.\n",
        "\n",
        "    Equivalent to `negative_sources_fn(num_negatives_per_target)`.\n",
        "    \"\"\"\n",
        "    return self.negative_sources_fn(self.num_negatives_per_target)\n",
        "\n",
        "  def generate_targets(self) -\u003e tf.Tensor:\n",
        "    \"\"\"Returns negative-sampled targets.\n",
        "\n",
        "    Equivalent to `negative_targets_fn(num_negatives_per_source)`.\n",
        "    \"\"\"\n",
        "    return self.negative_targets_fn(self.num_negatives_per_source)\n",
        "\n",
        "  @property\n",
        "  def total_negatives(self) -\u003e int:\n",
        "    return self.num_negatives_per_source + self.num_negatives_per_target\n",
        "\n",
        "\n",
        "def _provide_link_prediction_data_with_int_arithmetic_sampling(\n",
        "    graph_data: datasets.LinkPredictionGraphData,\n",
        "    source_sampling_spec: sampling_spec_pb2.SamplingSpec,\n",
        "    target_sampling_spec: sampling_spec_pb2.SamplingSpec\n",
        "    ) -\u003e tuple[_LinkPredictionSamplingProvider,\n",
        "               _LinkPredictionSamplingProvider]:\n",
        "  \"\"\"Dataset of subgraph samples for link prediction task.\"\"\"\n",
        "  adjacency = graph_data.edge_sets()[graph_data.target_edgeset].adjacency\n",
        "  validation_edges = graph_data.edge_split().validation_edges\n",
        "  negative_links_sampling = NegativeLinksSampling(\n",
        "      negative_sources_fn=functools.partial(\n",
        "          _uniform_random_int_negatives, output_dtype=adjacency.source.dtype,\n",
        "          max_node_id=graph_data.num_source_nodes - 1),\n",
        "      negative_targets_fn=functools.partial(\n",
        "          _uniform_random_int_negatives, output_dtype=adjacency.target.dtype,\n",
        "          max_node_id=graph_data.num_target_nodes - 1),\n",
        "      num_negatives_per_source=2, num_negatives_per_target=2)\n",
        "\n",
        "  train_edges_dataset_fn = functools.partial(\n",
        "      _create_labeled_edges_dataset_from_eager_tensors,\n",
        "      adjacency.source, adjacency.target, shuffle=True,\n",
        "      negative_links_sampling=negative_links_sampling)\n",
        "\n",
        "  validation_edges_dataset_fn = functools.partial(\n",
        "      _create_labeled_edges_dataset_from_eager_tensors,\n",
        "      validation_edges[0], validation_edges[1],\n",
        "      negative_links_sampling=negative_links_sampling)\n",
        "\n",
        "  # Sampling model.\n",
        "  sampler = ia_sampler.GraphSampler(\n",
        "      graph_data, sampling_mode=ia_sampler.EdgeSampling.WITH_REPLACEMENT)\n",
        "\n",
        "  sampling_model = link_samplers.create_link_prediction_sampling_model(\n",
        "      graph_data.graph_schema(),\n",
        "      source_sampling_spec=source_sampling_spec,\n",
        "      target_sampling_spec=target_sampling_spec,\n",
        "      source_edge_sampler_factory=sampler.make_edge_sampler,\n",
        "      target_edge_sampler_factory=sampler.make_edge_sampler,\n",
        "      node_features_accessor_factory=(\n",
        "          graph_data.create_node_features_lookup_factory()))\n",
        "\n",
        "  return (\n",
        "      _LinkPredictionSamplingProvider(\n",
        "          adjacency.source.shape[0], train_edges_dataset_fn, sampling_model),\n",
        "      _LinkPredictionSamplingProvider(\n",
        "          validation_edges.shape[-1], validation_edges_dataset_fn,\n",
        "          sampling_model))\n",
        "\n",
        "\n",
        "class _LinkPredictionSamplingProvider(SizedDatasetProvider):\n",
        "  \"\"\"DatasetProvider for Link Prediction.\"\"\"\n",
        "\n",
        "  def __init__(\n",
        "      self,\n",
        "      cardinality: int,\n",
        "      dataset_fn: Callable[[tf.distribute.InputContext], tf.data.Dataset],\n",
        "      sampling_model: tf.keras.Model,\n",
        "      sampling_batch_size: int = 100):\n",
        "    \"\"\"Constructs DatasetProvider for Link Prediction.\n",
        "\n",
        "    Args:\n",
        "      cardinality: Total number of examples in (one iteration of) dataset.\n",
        "      dataset_fn: Callable returns dataset that with each element containing\n",
        "        `(source ID, target ID, label)`. The first two will be batched to invoke\n",
        "        `sampling_model`.\n",
        "      sampling_model: tf.keras.Model to map items in `dataset` to `GraphTensor`.\n",
        "      sampling_batch_size: Batch size to send to `sampling_model`.\n",
        "    \"\"\"\n",
        "    super().__init__()\n",
        "    self._cardinality = cardinality\n",
        "    self._dataset_fn = dataset_fn\n",
        "    self._sampling_model = sampling_model\n",
        "    self._batch_size = sampling_batch_size\n",
        "\n",
        "  @property\n",
        "  def cardinality(self) -\u003e int:\n",
        "    return self._cardinality\n",
        "\n",
        "  def sample_subgraphs(self, source_ids, target_ids, labels):\n",
        "    return self._sampling_model((source_ids, target_ids)), labels\n",
        "\n",
        "  def attach_label(self, graph, label):\n",
        "    return graph.replace_features(\n",
        "        node_sets={'_readout': {'label': tf.expand_dims(label, 0)}})\n",
        "\n",
        "  def get_dataset(self, context: tf.distribute.InputContext) -\u003e tf.data.Dataset:\n",
        "    return (\n",
        "        self._dataset_fn(context)\n",
        "        .batch(self._batch_size)\n",
        "        .map(self.sample_subgraphs)\n",
        "        .unbatch()\n",
        "        .map(self.attach_label))\n",
        "\n",
        "\n",
        "def _uniform_random_int_negatives(\n",
        "    num_negatives: int, *,\n",
        "    output_dtype: tf.DType, max_node_id: int) -\u003e tf.Tensor:\n",
        "  return tf.random.uniform(\n",
        "      [num_negatives], minval=0, maxval=max_node_id, dtype=output_dtype)\n",
        "\n",
        "\n",
        "def _sample_negative_edges(\n",
        "    source_node: tf.Tensor, target_node: tf.Tensor, *,\n",
        "    negative_links_sampling: NegativeLinksSampling) -\u003e tuple[\n",
        "        tf.Tensor, tf.Tensor, tf.Tensor]:\n",
        "  \"\"\"Returns positive edge `{source,target}_node`, sampled negatives and labels.\n",
        "\n",
        "  Args:\n",
        "    source_node: scalar `tf.Tensor` containing ID of source node.\n",
        "    target_node: scalar `tf.Tensor` containing ID of target node.\n",
        "    negative_links_sampling: Contains functions that generate negative source\n",
        "      and target node IDs.\n",
        "\n",
        "  Returns:\n",
        "    Three tensors `(sources, targets, labels)`, all are vectors of length\n",
        "    `1 + N.num_negatives_per_source + N.num_negatives_per_target`, with\n",
        "    `N == negative_links_sampling`. First entry corresponds to positive edge,\n",
        "    with `sources[0] == source_node`, `targets[0] == target_node`, and\n",
        "    `labels[0] == 1`. Remainder entries correspond to negative edges, with\n",
        "    `labels[1:] == 0`, and `targets[i]` contains either `target_node` or a\n",
        "    random negative returned from `N.negative_targets_fn`. In case, it was\n",
        "    `target_node`, then `sources[i]` must contain random negative from\n",
        "    `N.negative_sources_fn` else it must be `sources[i] == source_id`.\n",
        "  \"\"\"\n",
        "  source_node = tf.ensure_shape(source_node, [])\n",
        "  target_node = tf.ensure_shape(target_node, [])\n",
        "  negative_sources = negative_links_sampling.generate_sources()\n",
        "  negative_targets = negative_links_sampling.generate_targets()\n",
        "\n",
        "  all_sources = tf.concat([\n",
        "      tf.expand_dims(source_node, 0),\n",
        "      tf.fill([negative_links_sampling.num_negatives_per_source], source_node),\n",
        "      negative_sources,\n",
        "  ], axis=0)\n",
        "\n",
        "  all_targets = tf.concat([\n",
        "      tf.expand_dims(target_node, 0),\n",
        "      negative_targets,\n",
        "      tf.fill([negative_links_sampling.num_negatives_per_target], target_node),\n",
        "  ], axis=0)\n",
        "\n",
        "  all_labels = tf.concat([\n",
        "      tf.ones([1], dtype=tf.float32),\n",
        "      tf.zeros([negative_links_sampling.total_negatives], dtype=tf.float32),\n",
        "  ], 0)\n",
        "  return all_sources, all_targets, all_labels\n",
        "\n",
        "\n",
        "def amend_negative_edges(\n",
        "    positive_edges_dataset: tf.data.Dataset,\n",
        "    *,\n",
        "    negative_links_sampling: NegativeLinksSampling,\n",
        "    ) -\u003e tf.data.Dataset:\n",
        "  \"\"\"Creates link-prediction dataset from `sources`, `targets`, and negatives.\n",
        "\n",
        "  It must be that `source_nodes` and `target_nodes` are both vectors with equal\n",
        "  shape, with each `(source_nodes[i], target_nodes[i])` indicating a positive\n",
        "  edge.\n",
        "\n",
        "  For each positive edge, a total of\n",
        "  ```num_negatives_per_source + num_negatives_per_target```\n",
        "  negative examples will be sampled (uniformly at random).\n",
        "\n",
        "  Args:\n",
        "    positive_edges_dataset: Dataset containing pairs of positive edge endpoints.\n",
        "      Each element must be `(source_id, target_id)`.\n",
        "    negative_links_sampling: Contains functions that generate negative source\n",
        "      and target node IDs.\n",
        "\n",
        "  Returns:\n",
        "    Dataset with each element being 3-scalars: `(source_id, target_id, label)`.\n",
        "    The `label` will be `1` if `source_id, target_id` are from\n",
        "    `positive_edges_dataset`. Otherwise, `label == 0`, in which case, either\n",
        "    `source_id` or `target_id` (but not both) are, respectively, determined by\n",
        "    `negative_sources_fn` or `negative_targets_fn` (and the other comes from\n",
        "    `positive_edges_dataset`).\n",
        "  \"\"\"\n",
        "  return (\n",
        "      positive_edges_dataset\n",
        "      .map(functools.partial(\n",
        "          _sample_negative_edges,\n",
        "          negative_links_sampling=negative_links_sampling))\n",
        "      .unbatch())\n",
        "\n",
        "\n",
        "## \u003cHelper functions for both tasks\u003e\n",
        "\n",
        "\n",
        "def _resolve_graph_data(\n",
        "    graph_data_class: type[datasets.InMemoryGraphData],\n",
        "    dataset_name: None|str = None,\n",
        "    graph_data: None|datasets.InMemoryGraphData = None,\n",
        "    ) -\u003e datasets.InMemoryGraphData:\n",
        "  \"\"\"Returns `graph_data` or `get_in_memory_graph_data(dataset_name)`.\n",
        "\n",
        "  Exactly one of `dataset_name` or `graph_data` must be given. If `graph_data`\n",
        "  is given, then it must be an instance of `graph_data_class`. If `dataset_name`\n",
        "  is given, then it must correspond to an `InMemoryGraphData` with type\n",
        "  `graph_data_class`. If any if these conditions are violated, `ValueError` will\n",
        "  be raised.\n",
        "\n",
        "  Args:\n",
        "    graph_data_class: Type of instance that will be returned.\n",
        "    dataset_name: Mutually-exclusive with `graph_data`. If given, it must be\n",
        "      name of `InMemoryGraphData` with type `graph_data_class`.\n",
        "    graph_data: Mutually-exclusive with `dataset_name`. If given, it must be\n",
        "      an instance of `graph_data_class`.\n",
        "\n",
        "  Returns:\n",
        "    `InMemoryGraphData` with type `graph_data_class`.\n",
        "  \"\"\"\n",
        "  if int(dataset_name is None) + int(graph_data is None) != 1:\n",
        "    raise ValueError('Expecting exactly one of `dataset_name` or `graph_data`')\n",
        "  if dataset_name is not None:\n",
        "    graph_data = datasets.get_in_memory_graph_data(dataset_name)\n",
        "    if not isinstance(graph_data, graph_data_class):\n",
        "      raise ValueError(\n",
        "          f'dataset \"{dataset_name}\" corresponds to object of type '\n",
        "          f'{type(graph_data)}. Expecting {graph_data_class}.')\n",
        "  else:\n",
        "    if not isinstance(graph_data, graph_data_class):\n",
        "      raise ValueError(\n",
        "          f'`graph_data` {graph_data} corresponds to object of type '\n",
        "          f'{type(graph_data)}. Expecting {graph_data_class}.')\n",
        "  assert graph_data is not None\n",
        "  return graph_data\n",
        "\n",
        "\n",
        "def _resolve_sampling_spec(\n",
        "    graph_schema: graph_schema_pb2.GraphSchema,\n",
        "    seed_node_set_name: graph_constants.NodeSetName,\n",
        "    sampling: Sampling\n",
        "    ) -\u003e sampling_spec_pb2.SamplingSpec:\n",
        "  \"\"\"Returns `SamplingSpec` from `sampling.spec` or `make_sampling_spec_tree`.\n",
        "\n",
        "  If `sampling.spec` is given, then it will be returned. Else, int list\n",
        "  `steps = sampling.sample_sizes` will be read and returned `SamplingSpec`\n",
        "  will sample `steps[0]` from every node set that `seed_node_set_name`\n",
        "  connects to, and from every sampled node, `steps[1]` neighbors will be sampled\n",
        "  for every node-set, etc, until depth `len(steps)`.\n",
        "\n",
        "  Args:\n",
        "    graph_schema: Only accessed if `sampling.spec is None`.\n",
        "    seed_node_set_name: Only accessed if `sampling.spec is None`. If accessed,\n",
        "      the result `SamplingSpec` will have its `seed_op.node_set_name` set to\n",
        "      this.\n",
        "    sampling: Must have either attribute `spec` or `sample_sizes` set.\n",
        "\n",
        "  Returns:\n",
        "    `sampling.spec` if not None. Else,\n",
        "    `make_sampling_spec_tree(graph_schema, seed_node_set_name,\n",
        "    sampling.sample_sizes)`. In which case, if `sampling.sample_sizes` is not\n",
        "    set, it will default to `Sampling.SMALL_SAMPLING`.\n",
        "  \"\"\"\n",
        "  if sampling is None:\n",
        "    sampling = Sampling.SMALL_SAMPLING\n",
        "\n",
        "  sampling_spec = sampling.spec\n",
        "  if sampling_spec is None:  # if `sampling`` is constructed `from_sizes`.\n",
        "    sampling_spec = sampling_spec_builder.make_sampling_spec_tree(\n",
        "        graph_schema, seed_node_set_name,\n",
        "        sample_sizes=list(sampling.sample_sizes))\n",
        "  return sampling_spec\n",
        "\n",
        "\n",
        "## \u003cTemporary HACK to be migrated to datasets.py or subgraph_pipeline.py\u003e\n",
        "\n",
        "\n",
        "# Helper functions to obtain FeatureAccessor that keeps '#id' feature.\n",
        "# NOTE: These can be copied to update `datasets.py`.\n",
        "def create_node_features_lookup_factory(\n",
        "    self_gd: datasets.InMemoryGraphData,\n",
        "    keep_id_feature: bool = False) -\u003e datasets.NodeFeaturesLookupFactory:\n",
        "  return functools.partial(\n",
        "      _node_features_lookup, node_sets=dict(self_gd.node_sets()), cache={},\n",
        "      resource_prefix=self_gd.name, keep_id_feature=keep_id_feature)\n",
        "\n",
        "\n",
        "class _Accessor(\n",
        "    tf.keras.layers.Layer, datasets.interfaces.KeyToFeaturesAccessor):\n",
        "  \"\"\"Wraps `NodeSet` with `call` that can select features for node subsets.\"\"\"\n",
        "\n",
        "  def __init__(\n",
        "      self, node_set: Any,\n",
        "      resource_name: str, keep_id_feature: bool):\n",
        "    super().__init__()\n",
        "    self._node_set = node_set\n",
        "    self._resource_name = resource_name\n",
        "    self._keep_id_feature = keep_id_feature\n",
        "\n",
        "  def call(self, keys: tf.RaggedTensor) -\u003e datasets.interfaces.Features:\n",
        "    \"\"\"Gathers features corresponding to (tf.int) node keys.\"\"\"\n",
        "    return {feature_name: tf.gather(feature_value, keys)\n",
        "            for feature_name, feature_value in self._node_set.features.items()\n",
        "            if self._keep_id_feature or feature_name != '#id'}\n",
        "\n",
        "  @property\n",
        "  def resource_name(self) -\u003e str:\n",
        "    return self._resource_name\n",
        "\n",
        "\n",
        "def _node_features_lookup(\n",
        "    node_set_name: graph_constants.NodeSetName,\n",
        "    *,\n",
        "    node_sets: dict[graph_constants.NodeSetName, Any],\n",
        "    cache: dict[graph_constants.NodeSetName, _Accessor],\n",
        "    resource_prefix: str,\n",
        "    keep_id_feature: bool) -\u003e datasets.interfaces.KeyToFeaturesAccessor:\n",
        "  \"\"\"Returns `KeyToFeaturesAccessor` for a given node set.\"\"\"\n",
        "  if node_set_name in cache:\n",
        "    return cache[node_set_name]\n",
        "  cache[node_set_name] = _Accessor(\n",
        "      node_sets[node_set_name], f'{resource_prefix}/nodes/{node_set_name}',\n",
        "      keep_id_feature)\n",
        "  return cache[node_set_name]\n",
        "\n",
        "\n",
        "class DataProviders:\n",
        "  pass\n",
        "\n",
        "data_providers = DataProviders()\n",
        "data_providers.provide_node_classification_data = provide_node_classification_data\n",
        "data_providers.provide_link_prediction_data = provide_link_prediction_data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KvDLw7UHzWzP"
      },
      "source": [
        "## Model\n",
        "\n",
        "in Tf-gnn we have multiple GNN layers. For this example, let's stick with GCN (Kipf and Welling, ICLR'17)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "abcFW423IFIc"
      },
      "outputs": [],
      "source": [
        "class GCNLayer(tf.keras.layers.Layer):\n",
        "\n",
        "  def __init__(self, out_dim=64):\n",
        "    super().__init__()\n",
        "    self._layer = gcn.GCNConv(out_dim)\n",
        "\n",
        "  def call(self, graph):\n",
        "    x = self._layer(graph, edge_set_name='edges')\n",
        "    return graph.replace_features(node_sets={'nodes': {'hidden_state': x}})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1FfGAt4PzdGm"
      },
      "outputs": [],
      "source": [
        "def model_fn(gtspec: tfgnn.GraphTensorSpec) -\u003e tf.keras.Model:\n",
        "  def init_node_state(node_set, node_set_name):\n",
        "    del node_set_name\n",
        "    return node_set['feat']\n",
        "\n",
        "  return tf.keras.Sequential([\n",
        "      tfgnn.keras.layers.MapFeatures(node_sets_fn=init_node_state),\n",
        "      GCNLayer(32),\n",
        "      GCNLayer(32),\n",
        "  ])\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_ujt7iOvH0R_"
      },
      "source": [
        "## Dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pgjcdVFYYmJh"
      },
      "source": [
        "To wrap your in-memory dataset with TF-GNN, you can inherit `InMemoryGraphData` (defined in\n",
        " https://github.com/tensorflow/gnn/blob/main/tensorflow_gnn/experimental/in_memory/datasets.py). In addition, you may inherit one of `NodeClassificationGraphData` or `LinkPredictionGraphData`.\n",
        "\n",
        " Inheriting `InMemoryGraphData` implies that you should define 3 methods:\n",
        "\n",
        "  * `node_counts()`: should return `dict` from node-set name to number of nodes.\n",
        "  * `node_features_dict()`: should return `dict` from node-set name to `dict` of feature name to feature `tf.Tensor`.\n",
        "  * `edge_lists()`: should return `dict` from `tuple` (source node-set name, edge-set name, target node-set name), e.g., `('author', 'writes', 'paper')`, to `tf.Tensor` of `[source_ids, target_ids]` where each `_ids` is an integer vector.\n",
        "\n",
        "In addition, implementing `NodeClassificationGraphData` implies that you should implement: `num_classes()`, `node_split()`, `labels()`, `labeled_nodeset`.\n",
        "\n",
        "Implementing `LinkPredictionGraphData` implies that you should implement: `edge_split()`, `target_edgeset()`.\n",
        "\n",
        "**Homework**: Write a class that accepts **any** `InMemoryGraphData` and converts it into `LinkPredictionGraphData` by randomly partitioning edges among `train:test:validate`. Implemented next."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "L8FKoDqFHzvZ"
      },
      "outputs": [],
      "source": [
        "class LinkPredictionWrapper(datasets.LinkPredictionGraphData):\n",
        "\n",
        "  def __init__(self, gd: datasets.InMemoryGraphData, target_edge_set_name: str):\n",
        "    super().__init__()\n",
        "    self._all_edge_lists = gd.edge_lists()\n",
        "    self._train_edge_lists = None  # Populated by _partition_target_edge_list.\n",
        "    self._edge_split = None  # Populated by _partition_target_edge_list.\n",
        "    assert target_edge_set_name in [k[1] for k in self._all_edge_lists.keys()]\n",
        "    self._gd = gd\n",
        "    self._target_edge_set_name = target_edge_set_name\n",
        "    self._partition_target_edge_list()\n",
        "\n",
        "  def _partition_target_edge_list(self):\n",
        "    self._train_edge_lists = {}\n",
        "    for edge_list_key, edges_tensor in self._all_edge_lists.items():\n",
        "      if edge_list_key[1] == self._target_edge_set_name:\n",
        "        # Number of edges per partition train:test:validate.\n",
        "        num_train_edges = int(0.8*edges_tensor.shape[1])\n",
        "        num_validate_edges = int(0.1*edges_tensor.shape[1])\n",
        "        num_test_edges = edges_tensor.shape[1] - num_train_edges - num_validate_edges\n",
        "\n",
        "        # Partition edges\n",
        "        random_state = np.random.RandomState(seed=1234)\n",
        "        perm = random_state.permutation(edges_tensor.shape[1])\n",
        "        train_edges = tf.gather(edges_tensor, perm[:num_train_edges], axis=1)\n",
        "        validate_edges = tf.gather(edges_tensor, perm[num_train_edges:num_train_edges + num_validate_edges], axis=1)\n",
        "        test_edges = tf.gather(edges_tensor, perm[num_train_edges + num_validate_edges:num_test_edges], axis=1)\n",
        "\n",
        "        self._train_edge_lists[edge_list_key] = train_edges\n",
        "        self._edge_split = datasets.EdgeSplit(\n",
        "            train_edges=train_edges,\n",
        "            validation_edges=validate_edges,\n",
        "            test_edges=test_edges,\n",
        "            negative_validation_edges=tf.zeros([2, 0], dtype=test_edges.dtype),\n",
        "            negative_test_edges=tf.zeros([2, 0], dtype=test_edges.dtype))\n",
        "      else:\n",
        "        self._train_edge_lists[edge_list_key] = edges_tensor\n",
        "\n",
        "  def node_counts(self):\n",
        "    return self._gd.node_counts()\n",
        "\n",
        "  def node_features_dicts(self):\n",
        "    return self._gd.node_features_dicts()\n",
        "\n",
        "  def edge_lists(self):\n",
        "    return self._train_edge_lists\n",
        "\n",
        "  def edge_split(self):\n",
        "    return self._edge_split\n",
        "\n",
        "  @property\n",
        "  def target_edgeset(self):\n",
        "    return self._target_edge_set_name\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vwjWNNVxyrTQ"
      },
      "source": [
        "The above class `LinkPredictionWrapper` allows us to wrap any graph dataset (e.g., node-classification) into a LinkPrediction dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DK-V2TSYy3Ql"
      },
      "outputs": [],
      "source": [
        "gd = datasets.get_in_memory_graph_data('cora')\n",
        "link_prediction_gd = LinkPredictionWrapper(gd, 'edges')\n",
        "\n",
        "tasked_data = data_providers.provide_link_prediction_data(\n",
        "    graph_data=link_prediction_gd)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4xSPnHuOH1ne"
      },
      "source": [
        "## Training Loop"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "s-R-0LLJzemN"
      },
      "outputs": [],
      "source": [
        "train_provider = tasked_data.provider\n",
        "validation_provider = tasked_data.validation_provider or train_provider\n",
        "\n",
        "global_batch_size = 128\n",
        "\n",
        "trainer = runner.KerasTrainer(\n",
        "    strategy=tf.distribute.MirroredStrategy(),\n",
        "    model_dir=runner.incrementing_model_dir(os.path.expanduser(\"~/tmp\")),\n",
        "    steps_per_epoch=train_provider.cardinality // global_batch_size,\n",
        "    validation_steps=min(100, validation_provider.cardinality // global_batch_size),\n",
        "    verbose=1,\n",
        ")\n",
        "\n",
        "runner.run(\n",
        "    train_ds_provider=train_provider,\n",
        "    valid_ds_provider=validation_provider,\n",
        "    model_fn=model_fn,\n",
        "    optimizer_fn=tf.keras.optimizers.Adam,\n",
        "    epochs=10,\n",
        "    trainer=trainer,\n",
        "    task=tasked_data.task,\n",
        "    gtspec=train_provider.get_dataset(\n",
        "        tf.distribute.InputContext()).element_spec,\n",
        "    global_batch_size=global_batch_size,\n",
        "    feature_processors=[\n",
        "        # Featurizers go here. (e.g., distance encoding).\n",
        "    ])"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "SuBd-HNGAq63",
        "KvDLw7UHzWzP"
      ],
      "last_runtime": {
        "build_target": "",
        "kind": "local"
      },
      "private_outputs": true,
      "provenance": [
        {
          "file_id": "1OfrwN-2DXTzbTk65xCu4kyCm6UOHpRvg",
          "timestamp": 1691320220866
        },
        {
          "file_id": "1gfUqjajkwWbqv3hYHb9Tgslja8Uab4Lb",
          "timestamp": 1689970193575
        }
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
